# -*- coding: utf-8 -*-
# created on 2017/05/16
# 求含参数函数零点个数 - Mathematica 求解法 + 特殊值法

from collections import defaultdict
from operator import itemgetter

from sympy import Union, FiniteSet, S, oo, latex, Interval

from mathsolver.functions.base import BaseValue, BaseVariable, BaseVariables, BaseFunction, BaseEq
from mathsolver.functions.daoshu.daoshuyunsuan import HanShuDaoShuYunSuan
from mathsolver.functions.hanshu.helper import (cal_canshu_from_expr, flatten, span_values_to_intervals, 
                                                get_rand_from_interval, my_groupby, check_func, check_inter)
from mathsolver.functions.mathematica.mathematicaSolve import MathematicaSolve


class LingDianGeShuHanCanT(BaseFunction):
    """含参数函数的零点 -  特殊值法

    输入：args[0] BaseEq 或者 BaseFunc
        args[1] optional 区间
    """
    canshu = None

    def solver(self, *args):
        func = check_func(args[0])
        expr, var = func.expression, func.var
        canshu = cal_canshu_from_expr(expr)
        self.canshu = canshu
        has_qujian, qujian = self.cal_qujian(args, func)

        # 求极值点
        deriv_func = HanShuDaoShuYunSuan().solver(func).output[0]
        if has_qujian:
            deriv_solver = MathematicaSolve().solver(BaseEq([deriv_func.expression, 0]),
                                                     BaseVariables([canshu, var]), {var: qujian}).output[0]
        else:
            deriv_solver = MathematicaSolve().solver(BaseEq([deriv_func.expression, 0]),
                                                     BaseVariables([canshu, var])).output[0]
        # print deriv_solver.value

        var_values = []
        canshu_ranges = []
        # 解结果为 {x: {1}, (x, a): {(log(-2*a), (-oo, 0))}}，或者 {x: {0, 2}}
        for k, v in deriv_solver.value.items():
            if isinstance(k, tuple):
                for v_piece in v:
                    k_v_dict = dict(zip(k, v_piece))
                    canshu_ranges.append(k_v_dict[canshu])
                    var_values.append(k_v_dict[var])
            else:
                if k == var:
                    var_values.append(v)
                else:
                    canshu_ranges.append(v)

        canshu_teshudian = []
        # print var_values, canshu_ranges
        if var_values:
            var_values = set(flatten(var_values))
            # 对每个极值点 x 求解 f(x) = 0，如果解不出来，跳过；如果解出来加到特殊点集合
            for var_value in var_values:
                expr_new = expr.subs(var, var_value)
                var_solver = MathematicaSolve().solver(BaseEq([expr_new, 0])).output[0].value.get(canshu)

                if var_solver:
                    canshu_teshudian.append(var_solver)

        # 抽取出每个参数区间的端点值，放到特殊点集合里面
        if canshu_ranges:
            canshu_ranges = flatten(canshu_ranges)
            for canshu_range in canshu_ranges:
                if isinstance(canshu_range, Interval):
                    canshu_teshudian.extend(canshu_range.args[:2])
                elif isinstance(canshu_range, Union):
                    for inter in canshu_range.args:
                        canshu_teshudian.extend(inter.args[:2])
                elif isinstance(canshu_range, FiniteSet):
                    canshu_teshudian.extend(list(canshu_range))
                else:
                    raise ValueError("unknown type %s" % type(canshu_range))

        canshu_teshudian = {item for item in flatten(canshu_teshudian) if item not in [oo, -oo]}

        # 按照 a 的特殊值，把参数区间划分为分段小区间
        canshu_intervals = span_values_to_intervals(canshu_teshudian)

        # 检测每个小区间上解的个数
        interval_solvecount = []
        for inter in canshu_intervals:
            canshu_rand = get_rand_from_interval(inter)
            expr_new = expr.subs(canshu, canshu_rand)
            # TODO: 可能会解不出来
            var_solver = MathematicaSolve().solver(BaseEq([expr_new, 0]), BaseVariable([var])).output[0].value
            # print '=> ', var_solver
            var_solver = var_solver.get(var)
            if var_solver:
                interval_solvecount.append([inter, len(var_solver)])
            else:
                interval_solvecount.append([inter, 0])

        # print interval_solvecount

        # 合并解的个数相同的区间
        group_by_count = my_groupby(interval_solvecount, itemgetter(1))
        result = {}
        self.steps.append(["", "根据函数的单调性和图像可知"])
        for ct, count_inter in group_by_count.items():
            interval = Union(*[inter for inter, _ in count_inter])
            result.update({ct: interval})
            self.steps.append(["", "当 %s ∈ %s 时，函数有 %s 个零点" % (latex(canshu), latex(interval), ct)])

        self.output.append(BaseValue(result))
        self.label.add("含参数函数的零点")
        return self

    def cal_qujian(self, args, func):
        # 计算给定区间
        qujian = S.Reals
        if len(args) == 2:
            qujian = check_inter(args[1])
        elif func.dingyiyu:
            qujian = check_inter(func.dingyiyu)
        else:
            given_def = self.search(func.var)
            if given_def:
                qujian = given_def
        has_qujian = qujian != S.Reals
        return has_qujian, qujian


class LingDianGeShuHanCanM(LingDianGeShuHanCanT):
    """含参数函数的零点 -  Mathematica 解
    
    输入：args[0] BaseEq 或者 BaseFunc
        args[1] optional 区间
    """
    canshu = None

    def solver(self, *args):
        func = check_func(args[0])
        expr, var = func.expression, func.var
        symbols = list(expr.free_symbols)
        canshu = cal_canshu_from_expr(expr)
        self.canshu = canshu
        # 计算给定区间
        has_qujian, qujian = self.cal_qujian(args, func)
        # 补充：区间上交点个数求参
        if has_qujian:
            math_solver = MathematicaSolve().solver(BaseEq([expr, 0]), BaseVariables(symbols),
                                                    {var: qujian}).output[0].value
        else:
            math_solver = MathematicaSolve().solver(BaseEq([expr, 0]), BaseVariables(symbols)).output[0].value

        vars_ = list(math_solver.keys())[0]  # 所有变量
        solutions = math_solver[vars_]
        index = vars_.index(canshu)  # variable 变量的 index

        # group by variable
        group_by_var = my_groupby(solutions, key=itemgetter(index))

        # 细节处理：如果是整数值，转换为 FiniteSet, 否则 Union 会报错c
        for k, v in group_by_var.items():
            if not isinstance(k, (Interval, Union)):
                k_fs = FiniteSet(k)
                group_by_var.update({k_fs: v})
                del group_by_var[k]

        # 细节处理：Mathematica 的解可能是冗余的，例如：{(0, 2): [(log(-b + 2)/log(2), (0, 2), 0)], [0, oo): [(log(b + 2)/log(2), [0, oo), 0)]})
        new_group_by_var = {}
        for k, v in group_by_var.items():
            for inter in group_by_var.keys():
                if k.issubset(inter) and k != inter:
                    new_group_by_var.update({k: v + group_by_var[inter]})
                else:
                    new_group_by_var.update({k: v})

        # group by count
        group_by_count = defaultdict(list)
        for k, v in new_group_by_var.items():
            ct = len(v)
            # 带入边界值进行检验
            if isinstance(k, Interval):
                inter_v = k.args[:2]
                bound = k.args[2:]
                new_bound = []
                for i_v, i_b in zip(inter_v, bound):
                    if not i_b:
                        temp = [a_v.subs({canshu: i_v}) for i, a_v in enumerate(v) if i != index]
                        if len(set(temp)) != ct:
                            group_by_count[len(set(temp))].append(FiniteSet(i_v))
                            new_bound.append(True)
                            continue

                    new_bound.append(i_b)

                k = Interval(*(inter_v + tuple(new_bound)))

            group_by_count[ct].append(k)

        # 合并区间，输出步骤
        result = {}
        self.steps.append(["", "根据函数的单调性和图像可知"])
        for ct, inter in group_by_count.items():
            interval = Union(*inter)
            result.update({ct: interval})
            self.steps.append(["", "当 %s ∈ %s 时，函数有 %s 个零点" % (latex(canshu), latex(interval), ct)])

        self.output.append(BaseValue(result))
        self.label.add("含参数函数交点个数")
        return self


class LingDianGeShuHanCan(BaseFunction):
    """含参数函数的零点 -  Mathematica + 特殊值法

    输入：args[0] BaseEq 或者 BaseFunc
        args[1] optional 区间
    """

    def solver(self, *args):
        clses = [LingDianGeShuHanCanM, LingDianGeShuHanCanT]
        for cls in clses:
            try:
                return cls(self.known).solver(*args)
            except Exception:
                continue

        raise RuntimeError("cannot solve problem")
